# 简单命令
import subprocess


def simple_demo():
    # 执行命令
    result = subprocess.run(['ls', '-lh'], capture_output=True, text=True)
    # 打印结果
    print(result.stdout)


def catch_error():
    """
    执行一个可能产生错误的命令，捕获错误信息

    :return:
    """
    # 执行命令
    result = subprocess.run(['ls', '-lh', '/home2'], capture_output=True, text=True)
    # 打印结果
    print("Output: ", result.stdout)
    print("Error: ", result.stderr)


def check_output():
    """
    只返回命令的输出
    :return:
    """
    output = subprocess.check_output(['ls', '-lh', '/home2'], text=True)
    print(output)


def return_code():
    """
    获取执行状态码
    :return:
    """
    result = subprocess.run(['unset', 'CUDA_VISIBLE_DEVICES'],shell=True, capture_output=True, text=True)
    # 打印结果
    print("Return code: ", result.returncode)
    print("Args: ", result.args)


def run_command(command, log_flag=False):
    """
    执行linux命令，并获取输出
    """
    # 使用 subprocess.Popen 执行命令
    process = subprocess.Popen(
        command,
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        bufsize=1,
        universal_newlines=True
    )

    # 实时读取输出
    if log_flag:
        with process.stdout:
            for line in iter(process.stdout.readline, ''):
                yield line

    # 等待命令执行完成
    process.wait()
    print(process.returncode)

    # 返回命令的退出状态
    return process.returncode


if __name__ == '__main__':
    # simple_demo()
    # catch_error()
    # check_output()
    return_code()
    # run_command("unset CUDA_VISIBLE_DEVICES")
